package com.own.stu.apm.agent;

import com.own.stu.apm.agent.core.plugin.AbstractClassEnhencePluginDefine;
import com.own.stu.apm.agent.core.plugin.EnhanceContext;
import com.own.stu.apm.agent.core.plugin.PluginBootstrap;
import com.own.stu.apm.agent.core.plugin.PluginFinder;
import lombok.extern.slf4j.Slf4j;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.dynamic.scaffold.TypeValidation;
import net.bytebuddy.utility.JavaModule;

import java.lang.instrument.Instrumentation;
import java.util.LinkedList;

@Slf4j
public class BaseAgent {
    public static void premain(String arg, Instrumentation instrumentation) {

        log.info("baseAgent.in");
        PluginFinder pluginFinder = null;
        try {
            pluginFinder = new PluginFinder(new PluginBootstrap().loadPlugins());
        } catch (Exception e) {
            log.error("pluginFinder error, init fail", e);
            return;
        }

        ByteBuddy byteBuddy = new ByteBuddy().with(TypeValidation.of(true));

        AgentBuilder.Default agentBuilder = new AgentBuilder.Default(byteBuddy);
        agentBuilder.type(pluginFinder.buildMatch())
                .transform(new CusTransformer(pluginFinder))
                .installOn(instrumentation);
    }


    private static class CusTransformer implements AgentBuilder.Transformer {
        PluginFinder pluginFinder;

        public CusTransformer(PluginFinder pluginFinder) {
            this.pluginFinder = pluginFinder;
        }

        @Override
        public DynamicType.Builder<?> transform(
                DynamicType.Builder<?> builder,
                TypeDescription typeDescription,
                ClassLoader classLoader,
                JavaModule module) {
            LinkedList<AbstractClassEnhencePluginDefine> pluginDefines =  pluginFinder.find(typeDescription);
            if(pluginDefines == null || pluginDefines.size() == 0){
                log.debug("match class: {}, not find plugin", typeDescription.getActualName());
                return builder;
            }

            DynamicType.Builder<?> newBuilder = builder;
            EnhanceContext enhanceContext = new EnhanceContext();
            for (AbstractClassEnhencePluginDefine pluginDefine : pluginDefines) {
                DynamicType.Builder<?> possibleBuilder = pluginDefine.define(typeDescription, newBuilder, classLoader, enhanceContext);
                if(possibleBuilder != null){
                    newBuilder = possibleBuilder;
                }
                if(enhanceContext.isEnhanced()){
                    log.debug("finish the enhench for {}", typeDescription.getTypeName());
                }
            }

            return newBuilder;
        }
    }
}
